﻿
Imports System.Threading
Imports std

Public Class ddoslite
    Public Const _queuelength As Int64 = 4096
    Public Const _rejectOverallPercent As Double = 0.6
    Public Const _rejectCurrentPercent As Double = 0.4
    Public Const _fakeInsertIntervalMS As Int64 = 512
    Public Const _resourceCount As UInt64 = 64
    Public Shared ReadOnly fakeInsertString As String = _
        ":_""/\<>[]{}?&^%" + createGUID() + createGUID() + ":_""/\<>[]{}?&^%"
    Private counterslock As sequentialLock = Nothing
    Private counters2lock As sequentialLock = Nothing

    Public Enum returnValue
        success
        fail
        accept
        reject
    End Enum

    Protected Class counter
        Private map As map(Of String, UInt64) = Nothing
        Private queue As queue(Of String) = Nothing
        Private queuelength As UInt64
        Private timer3Index As Guid = Nothing

        Public Sub New(ByVal queuelength As UInt64, ByVal fakeInsertIntervalMS As UInt64)
            map = New map(Of String, UInt64)()
            map(fakeInsertString) = queuelength
            queue = New queue(Of String)()
            Dim i As Int64
            For i = 0 To queuelength - 1
                queue.push(emptyString)
            Next
            Me.queuelength = queuelength

            'automatically insert something to queue, so do not need to care about time
            'to current constants, all queue will be cleaned in ~32 minutes
            timer3Index = timer3.addcall(AddressOf fakeInsert, fakeInsertIntervalMS)
        End Sub

        Private Sub fakeInsert()
            insert(fakeInsertString)
        End Sub

        Public Function insert(ByVal s As String) As UInt64
            Dim rtn As UInt64
            SyncLock queue
                debugassert(queue.size() = queuelength, "queue is not filled.")
                queue.push(s)
                Dim it As map(Of String, UInt64).iterator = Nothing
                it = map.find(s)
                If it = map.end() Then
                    it = map.insert(s, 1)
                    rtn = 1
                Else
                    With +it
                        .second += 1
                        rtn = .second
                    End With
                End If

                Dim s2 As String = Nothing
                s2 = queue.front()
                queue.pop()

                it = map.find(s2)
                If it <> map.end() Then
                    With +it
                        .second -= 1
                        If .second = 0 Then
                            debugassert(map.erase(it), "cannot erase " + s2 + " from map.")
                        End If
                    End With
                End If
            End SyncLock

            Return rtn
        End Function

        Protected Overrides Sub Finalize()
            debugassert(timer3.removecall(timer3Index), _
                        "failed to remove timer3 call for timer3Index " + timer3Index.ToString())
            MyBase.Finalize()
        End Sub
    End Class

    'count the percent of usage, return the percent
    Private Class counter2
        Private counter As map(Of String, UInt64) = Nothing
        Private total As UInt64

        Public Sub New(ByVal resouceCount As UInt64)
            counter = New map(Of String, UInt64)()
            total = resouceCount
            If Not debugassert(total > 0, "should not have 0 resources to use.") Then
                total = 1
            End If
        End Sub

        Public Function insert(ByVal s As String) As Double
            Dim count As UInt64
            SyncLock counter
                Dim it As map(Of String, UInt64).iterator = Nothing
                it = counter.find(s)
                If it = counter.end() Then
                    it = counter.insert(s, 1)
                Else
                    With (+it)
                        .second += 1
                    End With
                End If
                count = (+it).second
            End SyncLock

            Return count / total
        End Function

        Public Sub release(ByVal s As String)
            SyncLock counter
                Dim it As map(Of String, UInt64).iterator = Nothing
                it = counter.find(s)
                If it = counter.end() Then
                    debugassert(False, "should not release a key which is not in counter2, key = " + s)
                    it = counter.insert(s, 1)
                End If
                With (+it)
                    .second -= 1
                    If .second = 0 Then
                        counter.erase(it)
                    End If
                End With
            End SyncLock
        End Sub
    End Class

    Private counters As map(Of String, counter) = Nothing
    Private counters2 As map(Of String, counter2) = Nothing

    Protected Overridable Function queuelength(ByVal key As String) As Int64
        Return _queuelength
    End Function

    Protected Overridable Function rejectCurrentPercent(ByVal key As String) As Double
        Return _rejectCurrentPercent
    End Function

    Protected Overridable Function rejectOverallPercent(ByVal key As String) As Double
        Return _rejectOverallPercent
    End Function

    Protected Overridable Function fakeInsertIntervalMS(ByVal key As String) As Int64
        Return _fakeInsertIntervalMS
    End Function

    Protected Overridable Function resourceCount() As UInt64
        Return _resourceCount
    End Function

    Protected Function registerAccessingImpl(ByVal key As String, ByVal s As String) As Boolean
        'cannot trace emptystring, so can only ignore them
        If Not isemptystring(s) Then
            Dim rtn As Boolean = True

            Dim count As Int64
            counterslock.wait()
            Dim it As map(Of String, counter).iterator = Nothing
            it = counters.find(key)
            If it = counters.end() Then
                it = counters.insert(key, New counter(queuelength(key), fakeInsertIntervalMS(key)))
            End If
            count = (+it).second.insert(s)
            counterslock.release()
            If count / queuelength(key) >= rejectOverallPercent(key) Then
                rtn = False
            End If

            Dim percent As Double
            counters2lock.wait()
            Dim it2 As map(Of String, counter2).iterator = Nothing
            it2 = counters2.find(key)
            If it2 = counters2.end() Then
                it2 = counters2.insert(key, New counter2(resourceCount()))
            End If
            percent = (+it2).second.insert(s)
            counters2lock.release()

            If percent >= rejectCurrentPercent(key) Then
                rtn = False
            End If

            Return rtn
        Else
            Return True
        End If
    End Function

    Protected Sub releaseAccessingImpl(ByVal key As String, ByVal s As String)
        If Not isemptystring(s) Then
            counters2lock.wait()
            Dim it As map(Of String, counter2).iterator = Nothing
            it = counters2.find(key)
            If it <> counters2.end() Then
                With (+it)
                    .second.release(s)
                End With
            Else
                raiseError("releaseAccessingImpl happened before registerAccessingImpl.", errorHandle.errorType.warning)
            End If
            counters2lock.release()
        End If
    End Sub

    Public Function registerAccessing(ByVal ParamArray input() As String) As returnValue
        If input Is Nothing Then
            Return returnValue.fail
        End If
        Dim accept As Boolean = True
        Dim i As Int64
        For i = 0 To input.Length() - 2 Step 2
            If Not registerAccessingImpl(input(i), input(i + 1)) Then
                accept = False
            End If
        Next

        If accept Then
            Return returnValue.accept
        Else
            Return returnValue.reject
        End If
    End Function

    Public Function releaseAccessing(ByVal ParamArray input() As String) As returnValue
        If input Is Nothing Then
            Return returnValue.fail
        End If

        Dim i As Int64
        For i = 0 To input.Length() - 2 Step 2
            releaseAccessingImpl(input(i), input(i + 1))
        Next

        Return returnValue.success
    End Function

    Public Sub New()
        counters = New map(Of String, counter)()
        counters2 = New map(Of String, counter2)()
        counterslock = New sequentialLock()
        counters2lock = New sequentialLock()
    End Sub
End Class
